import asyncio
import time
from collections import deque

import numpy as np
import sounddevice as sd
from mcp_tool import mcp_tool_func

from xiaozhi_sdk import XiaoZhiWebsocket
from xiaozhi_sdk.config import (
    INPUT_AUDIO_CHANNELS,
    INPUT_AUDIO_FRAME_DURATION,
    INPUT_AUDIO_SAMPLE_RATE,
)

input_audio: deque[bytes] = deque()
MAC_ADDR = "00:22:44:66:88:00"

is_play_audio = False


async def message_handler_callback(message):
    print("message received:", message)


async def assistant_audio_play(audio_queue):
    global is_play_audio
    # 创建一个持续播放的流
    stream = sd.OutputStream(samplerate=INPUT_AUDIO_SAMPLE_RATE, channels=INPUT_AUDIO_CHANNELS, dtype=np.int16)
    stream.start()
    last_time = None

    while True:

        if not audio_queue:
            await asyncio.sleep(0.01)
            if last_time and time.time() - last_time > 1:
                is_play_audio = False
            continue

        is_play_audio = True
        pcm_data = audio_queue.popleft()
        stream.write(pcm_data)
        last_time = time.time()


class Client:
    def __init__(self, mac_address):
        self.mac_address = mac_address
        self.xiaozhi = None

    async def start(self):
        self.xiaozhi = XiaoZhiWebsocket(message_handler_callback)
        await self.xiaozhi.set_mcp_tool(mcp_tool_func(self.xiaozhi))
        await self.xiaozhi.init_connection(self.mac_address, aec=False)
        asyncio.create_task(assistant_audio_play(self.xiaozhi.output_audio_queue))

    def callback_func(self, indata, frames, time, status):
        pcm = (indata.flatten() * 32767).astype(np.int16).tobytes()
        input_audio.append(pcm)

    async def process_audio(self):
        while True:
            if not input_audio:
                await asyncio.sleep(0.02)
                continue
            pcm = input_audio.popleft()
            if not is_play_audio:
                await self.xiaozhi.send_audio(pcm)


async def main():
    client = Client(MAC_ADDR)
    await client.start()
    await asyncio.sleep(2)
    blocksize = INPUT_AUDIO_SAMPLE_RATE * INPUT_AUDIO_FRAME_DURATION // 1000

    with sd.InputStream(
        callback=client.callback_func,
        channels=INPUT_AUDIO_CHANNELS,
        samplerate=INPUT_AUDIO_SAMPLE_RATE,
        blocksize=blocksize,
    ):
        print("Recording... Press Ctrl+C to stop.")
        await client.process_audio()  # 持续处理音频


if __name__ == "__main__":
    asyncio.run(main())
